import numpy as np
import networkx as nx
import trimesh


def build_graph(mesh):
    edges = mesh.edges_unique
    length = mesh.edges_unique_length
    g = nx.Graph()
    for edge, L in zip(edges, length):
        g.add_edge(*edge, length=L)
    return g


def compress(data, selectors):
    return [d for d, s in zip(data, selectors) if s]


def discrete_mean_curvature_measure_on_mesh(mesh, radius):
    """
    Return the discrete mean curvature measure of a sphere centered
    at a point as detailed in 'Restricted Delaunay triangulations and normal
    cycle', Cohen-Steiner and Morvan.
    Parameters
    ----------
    radius : float, the sphere radius
    Returns
    --------
    mean_curvature: (n,) float, discrete mean curvature measure.
    """

    # axis aligned bounds
    bounds = np.column_stack((mesh.vertices - radius,
                              mesh.vertices + radius))

    # line segments that intersect axis aligned bounding box
    candidates = [list(mesh.face_adjacency_tree.intersection(b))
                  for b in bounds]
    # filter the unreasonable candidates
    G = build_graph(mesh)
    length_list = [nx.single_source_dijkstra_path_length(G, idx, cutoff=radius * 1.0, weight='length') for idx in
                   range(len(mesh.vertices))]
    candidates_new = []
    for i, (idx, x_candidates) in enumerate(zip(range(len(mesh.vertices)), candidates)):
        endpoints_idx = mesh.face_adjacency_edges[x_candidates]
        mask = [item in length_list[idx].keys() for item in endpoints_idx[:, 1]]
        x_candidates_new = compress(x_candidates, mask)
        candidates_new.append(x_candidates_new)

    mean_curv = np.empty(len(mesh.vertices))
    for i, (x, x_candidates) in enumerate(zip(mesh.vertices, candidates_new)):
        endpoints = mesh.vertices[mesh.face_adjacency_edges[x_candidates]]

        lengths = line_ball_intersection(
            endpoints[:, 0],
            endpoints[:, 1],
            center=x,
            radius=radius)
        angles = mesh.face_adjacency_angles[x_candidates]
        signs = np.where(mesh.face_adjacency_convex[x_candidates], 1, -1)
        mean_curv[i] = (lengths * angles * signs).sum() / 2

    return mean_curv


def line_ball_intersection(start_points, end_points, center, radius):
    """
    Compute the length of the intersection of a line segment with a ball.
    Parameters
    ----------
    start_points : (n,3) float, list of points in space
    end_points   : (n,3) float, list of points in space
    center       : (3,) float, the sphere center
    radius       : float, the sphere radius
    Returns
    --------
    lengths: (n,) float, the lengths.
    """

    # We solve for the intersection of |x-c|**2 = r**2 and
    # x = o + dL. This yields
    # d = (-l.(o-c) +- sqrt[ l.(o-c)**2 - l.l((o-c).(o-c) - r^**2) ]) / l.l
    L = end_points - start_points
    oc = start_points - center  # o-c
    r = radius
    ldotl = np.einsum('ij, ij->i', L, L)  # l.l
    ldotoc = np.einsum('ij, ij->i', L, oc)  # l.(o-c)
    ocdotoc = np.einsum('ij, ij->i', oc, oc)  # (o-c).(o-c)
    discrims = ldotoc ** 2 - ldotl * (ocdotoc - r ** 2)

    # If discriminant is non-positive, then we have zero length
    lengths = np.zeros(len(start_points))
    # Otherwise we solve for the solns with d2 > d1.
    m = discrims > 0  # mask
    d1 = (-ldotoc[m] - np.sqrt(discrims[m])) / ldotl[m]
    d2 = (-ldotoc[m] + np.sqrt(discrims[m])) / ldotl[m]

    # Line segment means we have 0 <= d <= 1
    d1 = np.clip(d1, 0, 1)
    d2 = np.clip(d2, 0, 1)

    # Length is |o + d2 l - o + d1 l|  = (d2 - d1) |l|
    lengths[m] = (d2 - d1) * np.sqrt(ldotl[m])

    return lengths


def sphere_ball_intersection(R, r):
    """
    Compute the surface area of the intersection of sphere of radius R centered
    at (0, 0, 0) with a ball of radius r centered at (R, 0, 0).
    Parameters
    ----------
    R : float, sphere radius
    r : float, ball radius
    Returns
    --------
    area: float, the surface are.
    """
    x = (2 * R ** 2 - r ** 2) / (2 * R)  # x coord of plane
    if x >= -R:
        return 2 * np.pi * R * (R - x)
    if x < -R:
        return 4 * np.pi * R ** 2


if __name__ == "__main__":

    # mesh = trimesh.load_mesh('~/finish_stl/073-a_lego_duplo.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/Margarine_800_tex.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/lm_models_obj_000011.stl')
    mesh = trimesh.load_mesh('~/finish_stl/Weight_24kg.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/040_large_marker.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/toy_rake.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/large_013_apple.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/011_banana.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/004_sugar_box.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/024_bowl.stl')
    # mesh = trimesh.load_mesh('~/finish_stl/hb_models_obj_000012.stl')
    # trimesh.smoothing.filter_humphrey(mesh)
    trimesh.smoothing.filter_humphrey(mesh)
    radii = [0.003]
    MeanCurvature = np.array(
        [discrete_mean_curvature_measure_on_mesh(mesh, r) / sphere_ball_intersection(1, r) for r in radii]).squeeze()
    threshold = 60
    mask_1 = MeanCurvature < threshold
    mask_2 = -threshold * 0.4 < MeanCurvature
    mask = mask_1 & mask_2
    mask_ = ~mask
    pc_pos = trimesh.PointCloud(mesh.vertices[mask_], colors=[255, 0, 0])
    pc_neg = trimesh.PointCloud(mesh.vertices[mask], colors=[0, 255, 0])
    scene = trimesh.Scene([mesh, pc_pos, pc_neg])
    scene.show()
